{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import matplotlib.pyplot as plt\n",
    "import polars as pl\n",
    "import datetime \n",
    "from tqdm import tqdm\n",
    "\n",
    "import plotly.express as px\n",
    "from plotly.subplots import make_subplots\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "from metric import score # Import event detection ap score function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These are variables to be used by the score function\n",
    "column_names = {\n",
    "    'series_id_column_name': 'series_id',\n",
    "    'time_column_name': 'step',\n",
    "    'event_column_name': 'event',\n",
    "    'score_column_name': 'score',\n",
    "}\n",
    "\n",
    "tolerances = {\n",
    "    'onset': [12, 36, 60, 90, 120, 150, 180, 240, 300, 360], \n",
    "    'wakeup': [12, 36, 60, 90, 120, 150, 180, 240, 300, 360]\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing data \n",
    "\n",
    "# Column transformations\n",
    "\n",
    "dt_transforms = [\n",
    "    pl.col('timestamp').str.to_datetime(), \n",
    "    (pl.col('timestamp').str.to_datetime().dt.year()-2000).cast(pl.UInt8).alias('year'), \n",
    "    pl.col('timestamp').str.to_datetime().dt.month().cast(pl.UInt8).alias('month'),\n",
    "    pl.col('timestamp').str.to_datetime().dt.day().cast(pl.UInt8).alias('day'), \n",
    "    pl.col('timestamp').str.to_datetime().dt.hour().cast(pl.UInt8).alias('hour')\n",
    "]\n",
    "\n",
    "data_transforms = [\n",
    "    pl.col('anglez').cast(pl.Int32), # Casting anglez to 16 bit integer\n",
    "    (pl.col('enmo')*1000).cast(pl.UInt32), # Convert enmo to 16 bit uint\n",
    "]\n",
    "\n",
    "train_series = pl.scan_parquet('train_series.parquet').with_columns(\n",
    "    dt_transforms + data_transforms\n",
    "    )\n",
    "\n",
    "train_events = pl.read_csv('train_events.csv').with_columns(\n",
    "    dt_transforms\n",
    "    ).drop_nulls()\n",
    "\n",
    "test_series = pl.scan_parquet('test_series.parquet').with_columns(\n",
    "    dt_transforms + data_transforms\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Removing null events and nights with mismatched counts from series_events\n",
    "mismatches = train_events.drop_nulls().group_by(['series_id', 'night']).agg([\n",
    "    ((pl.col('event') == 'onset').sum() == (pl.col('event') == 'wakeup').sum()).alias('balanced')\n",
    "    ]).sort(by=['series_id', 'night']).filter(~pl.col('balanced'))\n",
    "\n",
    "for mm in mismatches.to_numpy(): \n",
    "    train_events = train_events.filter(~((pl.col('series_id') == mm[0]) & (pl.col('night') == mm[1])))\n",
    "\n",
    "# Getting series ids as a list for convenience\n",
    "series_ids = train_events['series_id'].unique(maintain_order=True).to_list()\n",
    "\n",
    "# Updating train_series to only keep these series ids\n",
    "train_series = train_series.filter(pl.col('series_id').is_in(series_ids))\n",
    "\n",
    "features, feature_cols = [pl.col('hour')], ['hour']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Updating train_series to only keep these series ids\n",
    "train_series = train_series.filter(pl.col('series_id').is_in(series_ids))\n",
    "\n",
    "features, feature_cols = [pl.col('hour')], ['hour']\n",
    "\n",
    "for mins in [5, 30, 60*2, 60*8] :\n",
    "    \n",
    "    for var in ['enmo', 'anglez'] :\n",
    "        \n",
    "        features += [\n",
    "            pl.col(var).rolling_mean(12 * mins, center=True, min_periods=1).abs().cast(pl.UInt16).alias(f'{var}_{mins}m_mean'),\n",
    "            pl.col(var).rolling_max(12 * mins, center=True, min_periods=1).abs().cast(pl.UInt16).alias(f'{var}_{mins}m_max'),\n",
    "            pl.col(var).rolling_std(12 * mins, center=True, min_periods=1).abs().cast(pl.UInt16).alias(f'{var}_{mins}m_std')\n",
    "        ]\n",
    "\n",
    "        feature_cols += [ \n",
    "            f'{var}_{mins}m_mean', f'{var}_{mins}m_max', f'{var}_{mins}m_std'\n",
    "        ]\n",
    "\n",
    "        # Getting first variations\n",
    "        features += [\n",
    "            (pl.col(var).diff().abs().rolling_mean(12 * mins, center=True, min_periods=1)*10).abs().cast(pl.UInt32).alias(f'{var}_1v_{mins}m_mean'),\n",
    "            (pl.col(var).diff().abs().rolling_max(12 * mins, center=True, min_periods=1)*10).abs().cast(pl.UInt32).alias(f'{var}_1v_{mins}m_max'),\n",
    "            (pl.col(var).diff().abs().rolling_std(12 * mins, center=True, min_periods=1)*10).abs().cast(pl.UInt32).alias(f'{var}_1v_{mins}m_std')\n",
    "        ]\n",
    "\n",
    "        feature_cols += [ \n",
    "            f'{var}_1v_{mins}m_mean', f'{var}_1v_{mins}m_max', f'{var}_1v_{mins}m_std'\n",
    "        ]\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "id_cols = ['series_id', 'step', 'timestamp']\n",
    "\n",
    "train_series = train_series.with_columns(\n",
    "    features\n",
    ").select(id_cols + feature_cols)\n",
    "\n",
    "test_series = test_series.with_columns(\n",
    "    features\n",
    ").select(id_cols + feature_cols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_train_dataset(train_data, train_events, drop_nulls=False) :\n",
    "    \n",
    "    series_ids = train_data['series_id'].unique(maintain_order=True).to_list()\n",
    "    X, y = pl.DataFrame(), pl.DataFrame()\n",
    "    for idx in tqdm(series_ids) : \n",
    "        \n",
    "        # Normalizing sample features\n",
    "        sample = train_data.filter(pl.col('series_id')==idx).with_columns(\n",
    "            [(pl.col(col) / pl.col(col).std()).cast(pl.Float32) for col in feature_cols if col != 'hour']\n",
    "        )\n",
    "        \n",
    "        events = train_events.filter(pl.col('series_id')==idx)\n",
    "        \n",
    "        if drop_nulls : \n",
    "            # Removing datapoints on dates where no data was recorded\n",
    "            sample = sample.filter(\n",
    "                pl.col('timestamp').dt.date().is_in(events['timestamp'].dt.date())\n",
    "            )\n",
    "        \n",
    "        X = X.vstack(sample[id_cols + feature_cols])\n",
    "\n",
    "        onsets = events.filter((pl.col('event') == 'onset') & (pl.col('step').is_not_null()))['step'].to_list()\n",
    "        wakeups = events.filter((pl.col('event') == 'wakeup') & (pl.col('step').is_not_null()))['step'].to_list()\n",
    "\n",
    "        # NOTE: This will break if there are event series without any recorded onsets or wakeups\n",
    "        y = y.vstack(sample.with_columns(\n",
    "            sum([(onset <= pl.col('step')) & (pl.col('step') <= wakeup) for onset, wakeup in zip(onsets, wakeups)]).cast(pl.Boolean).alias('asleep')\n",
    "            ).select('asleep')\n",
    "            )\n",
    "    \n",
    "    y = y.to_numpy().ravel()\n",
    "    \n",
    "    return X, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_events(series, classifier) :\n",
    "    '''\n",
    "    Takes a time series and a classifier and returns a formatted submission dataframe.\n",
    "    '''\n",
    "    \n",
    "    series_ids = series['series_id'].unique(maintain_order=True).to_list()\n",
    "    events = pl.DataFrame(schema={'series_id':str, 'step':int, 'event':str, 'score':float})\n",
    "\n",
    "    for idx in tqdm(series_ids) : \n",
    "\n",
    "        # Collecting sample and normalizing features\n",
    "        scale_cols = [col for col in feature_cols if (col != 'hour') & (series[col].std() !=0)]\n",
    "        X = series.filter(pl.col('series_id') == idx).select(id_cols + feature_cols).with_columns(\n",
    "            [(pl.col(col) / series[col].std()).cast(pl.Float32) for col in scale_cols]\n",
    "        )\n",
    "\n",
    "        # Applying classifier to get predictions and scores\n",
    "        preds, probs = classifier.predict(X[feature_cols]), classifier.predict_proba(X[feature_cols])[:, 1]\n",
    "\n",
    "        #NOTE: Considered using rolling max to get sleep periods excluding <30 min interruptions, but ended up decreasing performance\n",
    "        X = X.with_columns(\n",
    "            pl.lit(preds).cast(pl.Int8).alias('prediction'), \n",
    "            pl.lit(probs).alias('probability')\n",
    "                        )\n",
    "        \n",
    "        # Getting predicted onset and wakeup time steps\n",
    "        pred_onsets = X.filter(X['prediction'].diff() > 0)['step'].to_list()\n",
    "        pred_wakeups = X.filter(X['prediction'].diff() < 0)['step'].to_list()\n",
    "        \n",
    "        if len(pred_onsets) > 0 : \n",
    "            \n",
    "            # Ensuring all predicted sleep periods begin and end\n",
    "            if min(pred_wakeups) < min(pred_onsets) : \n",
    "                pred_wakeups = pred_wakeups[1:]\n",
    "\n",
    "            if max(pred_onsets) > max(pred_wakeups) :\n",
    "                pred_onsets = pred_onsets[:-1]\n",
    "\n",
    "            # Keeping sleep periods longer than 30 minutes\n",
    "            sleep_periods = [(onset, wakeup) for onset, wakeup in zip(pred_onsets, pred_wakeups) if wakeup - onset >= 12 * 30]\n",
    "\n",
    "            for onset, wakeup in sleep_periods :\n",
    "                # Scoring using mean probability over period\n",
    "                score = X.filter((pl.col('step') >= onset) & (pl.col('step') <= wakeup))['probability'].mean()\n",
    "\n",
    "                # Adding sleep event to dataframe\n",
    "                events = events.vstack(pl.DataFrame().with_columns(\n",
    "                    pl.Series([idx, idx]).alias('series_id'), \n",
    "                    pl.Series([onset, wakeup]).alias('step'),\n",
    "                    pl.Series(['onset', 'wakeup']).alias('event'),\n",
    "                    pl.Series([score, score]).alias('score')\n",
    "                ))\n",
    "\n",
    "    # Adding row id column\n",
    "    events = events.to_pandas().reset_index().rename(columns={'index':'row_id'})\n",
    "\n",
    "    return events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "train_ids, val_ids = train_test_split(series_ids, train_size=0.7, random_state=42)\n",
    "\n",
    "# We will collect datapoints at 10 minute intervals for training for validating\n",
    "train_data = train_series.filter(pl.col('series_id').is_in(train_ids)).gather_every(12*5).collect()\n",
    "\n",
    "val_data = train_series.filter(pl.col('series_id').is_in(val_ids)).collect()\n",
    "val_solution = train_events.filter(pl.col('series_id').is_in(val_ids)).select(['series_id', 'event', 'step']).to_pandas()\n",
    "\n",
    "# Collecting datapoints at every 5 minutes\n",
    "train_data = train_series.filter(pl.col('series_id').is_in(series_ids)).gather_every(12 * 5).collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from dask_ml.wrappers import ParallelPostFit\n",
    "import joblib\n",
    "from dask.distributed import Client\n",
    "client=Client(\"172.16.120.173:8786\")\n",
    "\n",
    "X_train, y_train = make_train_dataset(train_data, train_events)\n",
    "del train_data\n",
    "\n",
    "# Training classifier\n",
    "rf_classifier = ParallelPostFit(RandomForestClassifier(n_estimators=500,\n",
    "                                    min_samples_leaf=25,\n",
    "                                    random_state=42,\n",
    "                                    n_jobs=-1)\n",
    ")\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "\n",
    "with joblib.parallel_backend('dask'):\n",
    "    rf_classifier.fit(X_train[feature_cols],y_train)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ray",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
